package com.gitee.ixtf.vertx.internal

import graphql.execution.instrumentation.InstrumentationState
import graphql.execution.instrumentation.SimplePerformantInstrumentation
import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters
import graphql.language.OperationDefinition.Operation.SUBSCRIPTION
import graphql.schema.DataFetcher
import graphql.schema.DataFetchingEnvironment
import io.vertx.core.Future
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono

object ToCompletionStageInstrumentation : SimplePerformantInstrumentation() {

  override fun instrumentDataFetcher(
      dataFetcher: DataFetcher<*>,
      parameters: InstrumentationFieldFetchParameters,
      state: InstrumentationState?
  ): DataFetcher<*> {
    if (parameters.isTrivialDataFetcher) return dataFetcher

    // Do not instrument data fetchers already instrumented by other Vert.x projects (e.g. Vert.x
    // RxJava)
    if (dataFetcher is InstrumentedDataFetcher) return dataFetcher

    if (parameters.executionContext.operationDefinition.operation == SUBSCRIPTION)
        return dataFetcher

    // parameters.environment.graphQlContext.put(InstrumentedDataFetcher::class.java, dataFetcher)
    return InstrumentedDataFetcher(dataFetcher)
  }

  private class InstrumentedDataFetcher(private val originalDataFetcher: DataFetcher<*>) :
      DataFetcher<Any> {
    override fun get(environment: DataFetchingEnvironment): Any? =
        when (val o: Any? = originalDataFetcher.get(environment)) {
          is Future<*> -> o.toCompletionStage()
          is Mono<*> -> o.toFuture()
          is Flux<*> -> o.collectList().toFuture()
          else -> o
        }
  }
}
